#include "kernel_operator.h"
#include "lib/matmul_intf.h"

using namespace AscendC;

class NLLLossSum {
 public:
  	__aicore__ inline NLLLossSum() {}
  	__aicore__ inline void Init(GM_ADDR x, GM_ADDR target, GM_ADDR weight,
                              GM_ADDR y ,
                              int32_t ignoreIndex, uint32_t sizeC,
                              uint32_t coreRowNum, uint32_t coreRowMod, TPipe* pipeIn) {
	this->pipe = pipeIn;

    this->ignoreIndex = ignoreIndex;
	this->sizeC = sizeC;
	this->coreNum = GetBlockNum();

	this->coreRowNum = coreRowNum;
	this->coreRowMod = coreRowMod;

	this->coreIndex = GetBlockIdx();
	this->beginRow = coreRowNum * this->coreIndex;
	this->beginRow += (coreIndex < coreRowMod) ? coreIndex : coreRowMod;
	this->coreRow = (coreIndex < coreRowMod) ? coreRowNum+1 : coreRowNum;


	yGm.SetGlobalBuffer((__gm__ DTYPE_Y*)y, sizeof(DTYPE_Y));
	// pipe->InitBuffer(ReduceXBuf, 32);

	if(this->coreIndex==0)
		InitGlobalMemory(yGm, 1, DTYPE_Y(0.0));

	// DataCopy(yGm, reduceXLocal, 8);
	// PipeBarrier<PIPE_ALL>();


	// InitGlobalMemory(yGm, 8, (float)0.0);
	xGm.SetGlobalBuffer((__gm__ DTYPE_X*)x + this->beginRow * this->sizeC,
                        this->coreRow * this->sizeC * sizeof(DTYPE_X));
    weightGm.SetGlobalBuffer((__gm__ DTYPE_WEIGHT*)weight,
                             sizeC * sizeof(DTYPE_WEIGHT));
	targetGm.SetGlobalBuffer(
        (__gm__ DTYPE_TARGET*)target + this->beginRow,
        this->coreRow * sizeof(DTYPE_TARGET));

    this->roundSizeC = ((this->sizeC+7) / 8) * 8;
    this->roundCoreRow = ((this->coreRow+7) / 8) * 8;

	pipe->InitBuffer(AllBuf, (12*roundCoreRow+roundSizeC) * sizeof(DTYPE_WEIGHT));

    // pipe->InitBuffer(AllWeightBuf, roundSizeC * sizeof(DTYPE_WEIGHT));

	// pipe->InitBuffer(AllTargetBuf, roundCoreRow * sizeof(DTYPE_TARGET));
	// // pipe.InitBuffer(AllXBuf, roundCoreRow * sizeof(DTYPE_TARGET));

	// // pipe.InitBuffer(inQueueTarget, BUFFER_NUM, ITER_N * sizeof(DTYPE_TARGET));

	// pipe->InitBuffer(GatherWeightBuf, roundCoreRow * sizeof(DTYPE_WEIGHT));
	// // this->gatherWeight = GatherWeightBuf.Get<DTYPE_WEIGHT>();

	// pipe->InitBuffer(InitIndexBuf, roundCoreRow * sizeof(DTYPE_TARGET));
	// pipe->InitBuffer(InitIndexXBuf, roundCoreRow * sizeof(DTYPE_TARGET));

	// // pipe.InitBuffer(IndexBuf, this->coreRow * sizeof(DTYPE_TARGET));

	
	// pipe->InitBuffer(GatherXBuf, roundCoreRow * sizeof(DTYPE_X));
	// pipe->InitBuffer(CopyInXBuf, roundCoreRow * sizeof(DTYPE_X) * 8);
	
	// pipe->InitBuffer(WorkBuf, (roundCoreRow ) * sizeof(DTYPE_X));


	}

	__aicore__ inline void process() {

		// printf("----sum class----\n coreNum=%d, coreIdx=%d, coreRow=%d \n",this->coreNum,this->coreIndex,this->coreRow);
		// PipeBarrier<PIPE_ALL>();
		// auto allTarget = AllTargetBuf.Get<DTYPE_TARGET>();
		// auto allWeight =AllWeightBuf.Get<DTYPE_WEIGHT>();
		// auto copyInX = CopyInXBuf.Get<DTYPE_X>();
		// auto gatherX = GatherXBuf.Get<DTYPE_X>();
		// auto indexLocal = InitIndexBuf.Get<DTYPE_TARGET>();
		// auto indexLocalUint = indexLocal.ReinterpretCast<uint32_t>();
		// auto indexXLocal = InitIndexXBuf.Get<DTYPE_TARGET>();
		// auto gatherWeight = GatherWeightBuf.Get<DTYPE_WEIGHT>();
		// auto reduceXLocal = ReduceXBuf.Get<DTYPE_X>();
		// auto workLocal = WorkBuf.Get<DTYPE_X>();
		// AscendC::DataCopyPadParams padParams{false, 0, 0, 0};
        // AscendC::DataCopyParams copyParams{1, 4, 0, 0};
		// AscendC::DataCopyExtParams extParams{1, 4, 0, 0, 0};

        auto allBuf = AllBuf.Get<DTYPE_WEIGHT>();
        auto copyInX = allBuf[0];
        auto indexLocal = allBuf[roundCoreRow*8].ReinterpretCast<DTYPE_TARGET>();
        auto indexXlocal = allBuf[roundCoreRow*9];
        auto indexWLocal = allBuf[roundCoreRow*10].ReinterpretCast<DTYPE_TARGET>();
        auto allTarget = allBuf[roundCoreRow*11].ReinterpretCast<DTYPE_TARGET>();
        auto allWeight = allBuf[roundCoreRow*12];
        auto indexLocalUint = indexLocal.ReinterpretCast<uint32_t>();
        auto indexWLocalUint = indexWLocal.ReinterpretCast<uint32_t>();

		AscendC::DataCopyPadParams padParams{false, 0, 0, 0};
		AscendC::DataCopyParams copyParams{2, 4, 0, 0};

		// CreateVecIndex(indexLocal, (int32_t)0, coreRow);
		// Muls(indexXLocal, indexLocal, (int32_t)sizeC , coreRow);
		// Muls(indexLocal, indexLocal, 32 , coreRow);


		/*
			copy in target
		*/

		DataCopy(allTarget, targetGm, roundCoreRow);
		// PipeBarrier<PIPE_ALL>();
        
		/*
			copy in weight
		*/

		DataCopy(allWeight, weightGm, roundSizeC);
        CreateVecIndex(indexLocal, (int32_t)0, coreRow);
        Muls(indexLocal, indexLocal, 32, coreRow);
		// PipeBarrier<PIPE_ALL>();

		/*
			gather x
		*/
		PipeBarrier<PIPE_ALL>();
		// Add(indexXLocal, indexXLocal, allTarget, coreRow);
		if (0 <= ignoreIndex && ignoreIndex < sizeC) {
			weightGm.SetValue(this->ignoreIndex, 0);
		  }
		Muls(indexWLocal, allTarget, 4, coreRow);
		PipeBarrier<PIPE_ALL>();
		// for (int i = 0; i < this->coreRow; i++) {
		// 	printf("[i,indexLocal,indexWLocal]= %d %d %d\n", i,indexLocal.GetValue(i),indexWLocal.GetValue(i));
		// }
		{
			int i = 0;
			// printf("before run 1 i=%d coreRow=%d \n",i,coreRow);
			const int loop1 = (int)coreRow - 8;
			const int loop2 = (int)coreRow - 2;
			while (i <= loop1) {
				// printf("run 1 i=%d loop1=%d \n",i,loop1);
				auto temp = allTarget.GetValue(i);
				// copyParams = {2, 4, (uint16_t)((allTarget.GetValue(i + 1) - temp + sizeC - 1)*4), 0};
				copyParams.srcStride = (uint16_t)((allTarget.GetValue(i + 1) - temp + sizeC - 1)*4);
				AscendC::DataCopyPad(copyInX[i * 8], xGm[i * sizeC + temp], copyParams, padParams);
				i+=2;
				temp = allTarget.GetValue(i);
				copyParams.srcStride = (uint16_t)((allTarget.GetValue(i + 1) - temp + sizeC - 1)*4);
				AscendC::DataCopyPad(copyInX[i * 8], xGm[i * sizeC + temp], copyParams, padParams);
				i+=2;
				temp = allTarget.GetValue(i);
				copyParams.srcStride = (uint16_t)((allTarget.GetValue(i + 1) - temp + sizeC - 1)*4);
				AscendC::DataCopyPad(copyInX[i * 8], xGm[i * sizeC + temp], copyParams, padParams);
				i+=2;
				temp = allTarget.GetValue(i);
				copyParams.srcStride = (uint16_t)((allTarget.GetValue(i + 1) - temp + sizeC - 1)*4);
				AscendC::DataCopyPad(copyInX[i * 8], xGm[i * sizeC + temp], copyParams, padParams);
				i+=2;
			}
			while (i <= loop2) {
				// printf("run 2\n");
				auto temp = allTarget.GetValue(i);
				// copyParams = {2, 4, (uint16_t)((allTarget.GetValue(i + 1) - temp + sizeC - 1)*4), 0};
				copyParams.srcStride = (uint16_t)((allTarget.GetValue(i + 1) - temp + sizeC - 1)*4);
				AscendC::DataCopyPad(copyInX[i * 8], xGm[i * sizeC + temp], copyParams, padParams);
				i+=2;
			}
			if (i < coreRow)
			{
				// printf("run 3\n");
				copyParams = {1, 4, 0, 0};
				AscendC::DataCopyPad(copyInX[i * 8], xGm[i * sizeC + allTarget.GetValue(i)], copyParams, padParams);
			}
		}

		// for (int i = 0; i < this->coreRow; i++) {
		// 	printf("[i,target,x]= %d %d %f\n", i,allTarget.GetValue(i), copyInX.GetValue(i*8));
		// }
		// PipeBarrier<PIPE_ALL>();
		Gather(allWeight, allWeight, indexWLocalUint, (uint32_t)0, this->coreRow);

		PipeBarrier<PIPE_ALL>();
		Gather(copyInX, copyInX, indexLocalUint, (uint32_t)0, this->coreRow);

		/*
			gather weight
		*/

		// PipeBarrier<PIPE_ALL>();

		/*
			Compute
		*/

		// for (int i = 0; i < this->coreRow; i++) {
		// 	printf("[i,target,weight,x]= %d %d %f %f\n", i,allTarget.GetValue(i),allWeight.GetValue(i), copyInX.GetValue(i));
		// }

		Mul(copyInX, copyInX, allWeight, this->coreRow);

		ReduceSum(copyInX, copyInX, copyInX, this->coreRow);
		// WholeReduceSum(copyInX, copyInX, this->coreRow, 1, 1, 1, 8);
		// BlockReduceSum<DTYPE_X>(copyInX, copyInX,1, this->coreRow, 1, 1, 8);

		// printf("reducexLocal: %f \n",copyInX.GetValue(0));
		float p = -1;
		Muls(copyInX, copyInX, p, 1);

		PipeBarrier<PIPE_ALL>();

		/*
			CopyToyGm
		*/
		AscendC::SetAtomicAdd<float>();
		AscendC::DataCopy(yGm, copyInX, 8);

	}


private:
	uint32_t sizeC;
	int32_t ignoreIndex;
	uint32_t coreRowNum,coreRowMod,coreNum,roundSizeC,roundCoreRow;
	uint32_t beginRow,coreIndex,coreRow;

	AscendC::TPipe* pipe;
	AscendC::GlobalTensor<DTYPE_X> xGm;
	AscendC::GlobalTensor<DTYPE_TARGET> targetGm;
	AscendC::GlobalTensor<DTYPE_WEIGHT> weightGm;
	AscendC::GlobalTensor<DTYPE_Y> yGm;

	// TBuf<QuePosition::VECCALC>  AllTargetBuf, AllWeightBuf;
	// TBuf<QuePosition::VECCALC> InitIndexBuf,InitIndexXBuf,IndexBuf,GatherWeightBuf,GatherXBuf, CopyInXBuf;
	// TBuf<QuePosition::VECCALC> ReduceXBuf,WorkBuf;
	// TBuf<QuePosition::VECCALC> FinalXBuf;

	// LocalTensor<DTYPE_WEIGHT> allWeight,gatherWeight;
	// LocalTensor<DTYPE_TARGET> allTarget;
	// LocalTensor<DTYPE_TARGET> initIndex;
	// LocalTensor<DTYPE_X> gatherX;

	// TBuf<QuePosition::VECCALC> AllTargetBuf, AllWeightBuf;
	// TBuf<QuePosition::VECCALC> InitIndexBuf, InitIndexXBuf, InitIndexWBuf, GatherXBuf, CopyInXBuf;
	// TBuf<QuePosition::VECCALC> ReduceXBuf;
    TBuf<QuePosition::VECCALC> AllBuf;
};